package com.abc.util;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;

import java.util.*;
import java.util.stream.Collectors;

/**
 * RRF (Reciprocal Rank Fusion) 算法工具类
 * 用于融合多个排序结果，提升搜索质量
 */
@Slf4j
public class RRFUtil {

    /**
     * 默认的 RRF 常数 k，用于平滑排名
     */
    private static final int DEFAULT_K = 60;

    /**
     * 使用 RRF 算法融合多个文档列表
     * 
     * @param documentLists 多个文档排序列表
     * @param k RRF 常数，默认为 60
     * @return 融合后的文档列表
     */
    public static List<Document> fuseRankings(List<List<Document>> documentLists, int k) {
        if (documentLists == null || documentLists.isEmpty()) {
            return Collections.emptyList();
        }

        // 移除空列表
        List<List<Document>> validLists = documentLists.stream()
                .filter(list -> list != null && !list.isEmpty())
                .collect(Collectors.toList());

        if (validLists.isEmpty()) {
            return Collections.emptyList();
        }

        if (validLists.size() == 1) {
            return new ArrayList<>(validLists.get(0));
        }

        // 计算每个文档的 RRF 分数
        Map<String, RRFScore> documentScores = new HashMap<>();

        for (List<Document> docList : validLists) {
            for (int rank = 0; rank < docList.size(); rank++) {
                Document doc = docList.get(rank);
                String docId = getDocumentId(doc);
                
                // RRF 分数计算: 1 / (k + rank + 1)
                double rrfScore = 1.0 / (k + rank + 1);
                
                documentScores.computeIfAbsent(docId, id -> new RRFScore(doc))
                        .addScore(rrfScore);
            }
        }

        // 按 RRF 分数排序并返回结果
        List<Document> fusedResults = documentScores.values().stream()
                .sorted((a, b) -> Double.compare(b.getTotalScore(), a.getTotalScore()))
                .map(RRFScore::getDocument)
                .collect(Collectors.toList());

        log.debug("RRF 融合完成: 输入列表数={}, 融合后文档数={}", validLists.size(), fusedResults.size());
        
        return fusedResults;
    }

    /**
     * 使用默认 k 值的 RRF 融合
     */
    public static List<Document> fuseRankings(List<List<Document>> documentLists) {
        return fuseRankings(documentLists, DEFAULT_K);
    }

    /**
     * 对单个文档列表进行多重排序并融合
     * 这里可以基于不同的相似度算法或评分标准进行多次排序
     * 
     * @param documents 原始文档列表
     * @param query 查询文本
     * @return RRF 融合后的文档列表
     */
    public static List<Document> enhanceRanking(List<Document> documents, String query) {
        if (documents == null || documents.isEmpty()) {
            return Collections.emptyList();
        }

        if (documents.size() == 1) {
            return new ArrayList<>(documents);
        }

        // 创建多个排序策略
        List<List<Document>> rankings = new ArrayList<>();
        
        // 1. 保持原始排序（基于向量相似度）
        rankings.add(new ArrayList<>(documents));
        
        // 2. 基于文档长度的排序（较长的文档可能包含更多信息）
        List<Document> lengthSorted = documents.stream()
                .sorted((a, b) -> Integer.compare(b.getText().length(), a.getText().length()))
                .collect(Collectors.toList());
        rankings.add(lengthSorted);
        
        // 3. 基于关键词匹配度的排序
        if (query != null && !query.trim().isEmpty()) {
            List<Document> keywordSorted = documents.stream()
                    .sorted((a, b) -> {
                        int scoreA = calculateKeywordScore(a.getText(), query);
                        int scoreB = calculateKeywordScore(b.getText(), query);
                        return Integer.compare(scoreB, scoreA);
                    })
                    .collect(Collectors.toList());
            rankings.add(keywordSorted);
        }

        // 使用 RRF 融合多个排序结果
        List<Document> fusedResults = fuseRankings(rankings);
        
        log.info("RRF 增强排序完成: 原始文档数={}, 排序策略数={}, 融合后文档数={}", 
                documents.size(), rankings.size(), fusedResults.size());
        
        return fusedResults;
    }

    /**
     * 计算文档文本与查询的关键词匹配分数
     */
    private static int calculateKeywordScore(String text, String query) {
        if (text == null || query == null) {
            return 0;
        }
        
        String lowerText = text.toLowerCase();
        String lowerQuery = query.toLowerCase();
        
        // 简单的关键词匹配计分
        String[] queryWords = lowerQuery.split("\\s+");
        int score = 0;
        
        for (String word : queryWords) {
            if (word.length() > 1) { // 忽略单字符
                int count = countOccurrences(lowerText, word);
                score += count * word.length(); // 较长的词权重更高
            }
        }
        
        return score;
    }

    /**
     * 计算子字符串在文本中的出现次数
     */
    private static int countOccurrences(String text, String substring) {
        int count = 0;
        int index = 0;
        while ((index = text.indexOf(substring, index)) != -1) {
            count++;
            index += substring.length();
        }
        return count;
    }

    /**
     * 获取文档的唯一标识符
     * 优先使用文档ID，如果没有则使用文本内容的哈希值
     */
    private static String getDocumentId(Document doc) {
        if (doc.getId() != null && !doc.getId().isEmpty()) {
            return doc.getId();
        }
        // 使用文本内容的哈希值作为唯一标识
        return String.valueOf(doc.getText().hashCode());
    }

    /**
     * RRF 分数累积器
     */
    private static class RRFScore {
        private final Document document;
        private double totalScore = 0.0;

        public RRFScore(Document document) {
            this.document = document;
        }

        public void addScore(double score) {
            this.totalScore += score;
        }

        public double getTotalScore() {
            return totalScore;
        }

        public Document getDocument() {
            return document;
        }
    }
}
